burst prefix-cache affinity for prompt-sharing requests#1727
Conversation
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
…roups Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
elevran
left a comment
There was a problem hiding this comment.
Thanks - the burst producer addresses a real problem in RL workoads and the measured results are compelling.
Some correctness issues need addressing before merge, plus several gaps in the batch lifecycle tests. All comments inlined.
A;so - there is no linked GitHub issue. Per project conventions, non-trivial work must be tracked in an issue before merge. Please follow project guidelines on future PRs.
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
There is a git issue but it wasn't linked. I just linked it #1726 |
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
Signed-off-by: Ezra Silvera <ezra@il.ibm.com>
|
@elevran Thanks for the detailed review. I address all issues you raised |
Summary
A burst of prefix-sharing prompts arriving in one window all score zero on the prefix-cache scorer (no cache is warm yet), so load balancing scatters them and the shared prefix is prefilled redundantly. This adds a request-level data producer that, within a batch window, assigns prompt-sharing requests jointly and before routing - co-locating those that share a prefix so it is prefilled once. It covers identical-prompt bursts (e.g. RL rollout samples) and distinct prompts that share a prefix (shared system prompt, RAG context, one prompt contained in another), bounded by a per-replica fair-share cap so co-location never unbalances the cluster.
Fixes: #1726
What changed
Two commits:
refactor(prefix): extract block-hashing into shared prefixhash package - behavior-preserving move of the prompt block-hashing (
GetBlockHashes,HashBlock, chained block hash) out ofapproximateprefixinto a newdataproducer/prefixhashpackage, so multiple prefix producers derive identical block hashes.approximateprefixkeeps atype blockHash = prefixhash.BlockHashalias; no behavior change.feat(prefix): add burst prefix cache producer for grouped requests - new
dataproducer/burstprefixpackage implementing aDataProducer:windowDurationMsare collected; when the window seals they are assigned jointly. Processing the batch under a lock removes the per-request race that makes near-simultaneous arrivals unschedulable together.minColocateBlocks > 0and it shares at least that many leading blocks with another request in the batch (so one prompt contained in another, or many prompts sharing a system preamble, co-locate); a request overlapping no other, and any prefix-less request, is scored zero everywhere so other scorers decide.maxPerReplica(k) before spilling to the next least-loaded replica;k = -1places the whole group on one replica. Identical groups are placed first and kept whole - same-prompt co-location is the firm structure and is never broken by an attaching singleton - then prefix-sharing units attach to it. A unit prefers a replica that already holds a unit sharing at leastminColocateBlocksleading blocks and is still under its fair share of the batch (total placed samples / replicas); otherwise placement is load-balanced. Units are placed longest-prefix first so shorter units match against the richest index. The fair-share cap is what makes co-location balance-safe: many units sharing one prefix spread across replicas up to their share rather than stampeding onto the seed replica.0keeps placement to identical groups, purely load-balanced.PrefixCacheMatchInfo(full match on the assigned replica, zero elsewhere) and reuses the existingprefix-cache-scorerunchanged - point itsprefixMatchInfoProducerNameat this producer. This is the same producer/scorer split already used by the precise and approximate prefix paths.Also: registration in
cmd/epp/runner/runner.go, a sample configdeploy/config/epp-burst-prefix-cache-config.yaml, package READMEs, and the producer index table entry.How to enable
Wire
token-producer -> burst-prefix-cache-producer -> prefix-cache-scorer(see the sample config). Key parameters:windowDurationMsmaxPerReplicablockSizeTokensmaxPrefixTokensToMatchminColocateBlocksThe scorer reads from the producer by name:
Testing
kspreads evenly and fills before spilling; a lone group of one and empty-prefix requests get no affinity; distinct groups spread across replicas. WithminColocateBlocksset: prefix-sharing families co-locate within their fair share; many groups sharing one prefix spread across replicas (no stampede); a shared prefix below the threshold does not co-locate; a single request overlapping a group attaches to it while the identical group stays whole; a lone request overlapping nothing gets no affinity; and overlapping singletons co-locate with no identical groups present (the shared-context case).Produce: a concurrent burst of identical prompts all co-locate on one replica; a lone request gets no affinity.make presubmitpasses (format, lint, typos, govulncheck, tags).Results (measured)
Measured on a real verl + llm-d integration: llm-d running in no-kube (local) mode, with EPP serving as the vLLM endpoint picker for verl's rollout requests. The only variable between the two arms is the burst producer; same GRPO/GSM8K rollout workload (4 vLLM replicas, TP=1, Qwen3-4B, group sampling G=8, 40 steps).
windowDurationMs=100,maxPerReplica=-1).Every full group co-locates onto one replica (mean 1.00 distinct replicas, down from 2.83), the prefix-cache hit rate rises 51.2% -> 68.3%, and mean rollout gen time per step drops 21.0 s -> 16.4 s (~22%). Latency does not regress - it improves, and most of the gain is in the tail: the per-step series flattens from spikes near 45 s to a steady ~13-24 s band. Co-locating a group does not serialize decode (vLLM continuous-batches the concurrent sequences) and removing redundant prefill frees headroom.
These timing numbers are a floor, not a ceiling. This workload is decode-dominant with very short prompts (~55-88 tokens, about one block), so prefill is a small fraction of end-to-end time and the latency win here mostly reflects freed prefill headroom rather than the prefill saving itself. As prompts grow - reasoning traces, long system prompts, RAG context - prefill becomes a far larger share of the work, so prefilling a shared prefix once instead of on every replica should yield substantially larger timing wins. The co-location and hit-rate gains shown here are the same mechanism; their latency payoff scales with prompt length.
Group co-location - groups go from 13.6% on a single replica (mean 2.83) to 100% (mean 1.00):
Prefix-cache hit rate - rises 51.2% -> 68.3%:
Mean rollout gen time per step - drops 21.0 s -> 16.4 s, tail spikes flattened: